package cc.shacocloud.mirage.loader;

import cc.shacocloud.mirage.loader.jar.Handler;
import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;

import java.io.IOException;
import java.net.JarURLConnection;
import java.net.URL;
import java.net.URLClassLoader;
import java.net.URLConnection;
import java.security.AccessController;
import java.security.PrivilegedExceptionAction;
import java.util.Enumeration;
import java.util.function.Supplier;
import java.util.jar.JarFile;
import java.util.jar.Manifest;

/**
 * mirage 启动类加载器
 *
 * @author 思追(shaco)
 */
public class MirageClassLoader extends URLClassLoader {
    
    static {
        // 注册为支持并行的
        ClassLoader.registerAsParallelCapable();
    }
    
    private final boolean exploded;
    
    private final cc.shacocloud.mirage.loader.jar.JarFile jarFile;
    
    private final Object packageLock = new Object();
    
    private volatile DefinePackageCallType definePackageCallType;
    
    public MirageClassLoader(URL[] urls, ClassLoader parent) {
        this(false, urls, parent);
    }
    
    public MirageClassLoader(boolean exploded, URL[] urls, ClassLoader parent) {
        this(exploded, null, urls, parent);
    }
    
    public MirageClassLoader(boolean exploded, cc.shacocloud.mirage.loader.jar.JarFile jarFile, URL[] urls, ClassLoader parent) {
        super(urls, parent);
        this.exploded = exploded;
        this.jarFile = jarFile;
    }
    
    @Override
    public URL findResource(String name) {
        if (this.exploded) {
            return super.findResource(name);
        }
        Handler.setUseFastConnectionExceptions(true);
        try {
            return super.findResource(name);
        } finally {
            Handler.setUseFastConnectionExceptions(false);
        }
    }
    
    @Override
    public Enumeration<URL> findResources(String name) throws IOException {
        if (this.exploded) {
            return super.findResources(name);
        }
        Handler.setUseFastConnectionExceptions(true);
        try {
            return new UseFastConnectionExceptionsEnumeration(super.findResources(name));
        } finally {
            Handler.setUseFastConnectionExceptions(false);
        }
    }
    
    @Override
    protected Class<?> loadClass(String name, boolean resolve) throws ClassNotFoundException {
        if (this.exploded) {
            return super.loadClass(name, resolve);
        }
        
        Handler.setUseFastConnectionExceptions(true);
        try {
            try {
                definePackageIfNecessary(name);
            } catch (IllegalArgumentException ex) {
                if (getDefinedPackage(name) == null) {
                    // 这不应该发生，因为 IllegalArgumentException 指示包已被定义，因此，getDefinedPackage(name) 不应返回空值
                    throw new AssertionError("包 " + name + " 已定义，但未能匹配");
                }
            }
            return super.loadClass(name, resolve);
        } finally {
            Handler.setUseFastConnectionExceptions(false);
        }
    }
    
    /**
     * 在进行 {@code findClass} 调用之前定义包，这对于确保嵌套 JAR 的相应清单与包相关联是必需的。
     */
    protected void definePackageIfNecessary(@NotNull String className) {
        int lastDot = className.lastIndexOf('.');
        if (lastDot >= 0) {
            String packageName = className.substring(0, lastDot);
            if (getDefinedPackage(packageName) == null) {
                try {
                    definePackage(className, packageName);
                } catch (IllegalArgumentException ex) {
                    if (getDefinedPackage(packageName) == null) {
                        // 这不应该发生，因为 IllegalArgumentException 指示包已被定义，因此，getDefinedPackage(name) 不应返回空值
                        throw new AssertionError("包 " + packageName + " 已定义，但未能匹配");
                    }
                }
            }
        }
    }
    
    private void definePackage(String className, String packageName) {
        try {
            AccessController.doPrivileged((PrivilegedExceptionAction<Object>) () -> {
                String packageEntryName = packageName.replace('.', '/') + "/";
                String classEntryName = className.replace('.', '/') + ".class";
                for (URL url : getURLs()) {
                    try {
                        URLConnection connection = url.openConnection();
                        if (connection instanceof JarURLConnection) {
                            JarFile jarFile = ((JarURLConnection) connection).getJarFile();
                            if (jarFile.getEntry(classEntryName) != null && jarFile.getEntry(packageEntryName) != null
                                    && jarFile.getManifest() != null) {
                                definePackage(packageName, jarFile.getManifest(), url);
                                return null;
                            }
                        }
                    } catch (IOException ignore) {
                    }
                }
                return null;
            }, AccessController.getContext());
        } catch (java.security.PrivilegedActionException ignore) {
        }
    }
    
    @Override
    protected Package definePackage(String name, Manifest man, URL url) throws IllegalArgumentException {
        if (!this.exploded) {
            return super.definePackage(name, man, url);
        }
        synchronized (this.packageLock) {
            return doDefinePackage(DefinePackageCallType.MANIFEST, () -> super.definePackage(name, man, url));
        }
    }
    
    @Override
    protected Package definePackage(String name, String specTitle, String specVersion, String specVendor,
                                    String implTitle, String implVersion, String implVendor, URL sealBase) throws IllegalArgumentException {
        if (!this.exploded) {
            return super.definePackage(name, specTitle, specVersion, specVendor, implTitle, implVersion, implVendor, sealBase);
        }
        
        synchronized (this.packageLock) {
            if (this.definePackageCallType == null) {
                Manifest manifest = getManifest(this.jarFile);
                if (manifest != null) {
                    return definePackage(name, manifest, sealBase);
                }
            }
            return doDefinePackage(DefinePackageCallType.ATTRIBUTES, () -> super.definePackage(name, specTitle,
                    specVersion, specVendor, implTitle, implVersion, implVendor, sealBase));
        }
    }
    
    @Nullable
    private Manifest getManifest(cc.shacocloud.mirage.loader.jar.JarFile jarFile) {
        try {
            return (jarFile != null) ? jarFile.getManifest() : null;
        } catch (IOException ex) {
            return null;
        }
    }
    
    private <T> T doDefinePackage(DefinePackageCallType type, @NotNull Supplier<T> call) {
        DefinePackageCallType existingType = this.definePackageCallType;
        try {
            this.definePackageCallType = type;
            return call.get();
        } finally {
            this.definePackageCallType = existingType;
        }
    }
    
    /**
     * 为定义包而进行的不同类型的调用。我们跟踪这些分解的 jar，以便我们可以检测应应用清单属性的包
     */
    private enum DefinePackageCallType {
        
        /**
         * 来自具有清单的资源的定义包调用
         */
        MANIFEST,
        
        /**
         * 具有一组直接属性的定义包调用
         */
        ATTRIBUTES
        
    }
    
    private static class UseFastConnectionExceptionsEnumeration implements Enumeration<URL> {
        
        private final Enumeration<URL> delegate;
        
        UseFastConnectionExceptionsEnumeration(Enumeration<URL> delegate) {
            this.delegate = delegate;
        }
        
        @Override
        public boolean hasMoreElements() {
            Handler.setUseFastConnectionExceptions(true);
            try {
                return this.delegate.hasMoreElements();
            } finally {
                Handler.setUseFastConnectionExceptions(false);
            }
            
        }
        
        @Override
        public URL nextElement() {
            Handler.setUseFastConnectionExceptions(true);
            try {
                return this.delegate.nextElement();
            } finally {
                Handler.setUseFastConnectionExceptions(false);
            }
        }
        
    }
    
}
